from functools import cached_property
from typing import Sequence
from .seq import MTSequence
from .gene import Transcript, Region, RegionType
from .snv import SNV, TranscriptAnnotation


class SNP(SNV):
    @cached_property
    def s_start(self) -> int:
        return self.pos

    @cached_property
    def s_ref(self) -> Sequence | MTSequence:
        return self.ref

    @cached_property
    def s_alt(self) -> Sequence | MTSequence:
        return self.alt

    def annotate_utr(self, transcript: Transcript, region: Region, utr_size: int) -> TranscriptAnnotation:
        distance = transcript.calc_utr_distance(region, utr_size, self.pos)
        trans_anno = TranscriptAnnotation.from_transcript(transcript, region)
        ref, alt = (self.s_ref, self.s_alt) if transcript.strand == "+" else (self.s_ref.reverse_complement(), self.s_alt.reverse_complement())
        trans_anno.na_change = f"c.{distance}{ref}>{alt}"
        return trans_anno

    def annotate_intron(self, transcript: Transcript, region: Region, cds_size: int) -> TranscriptAnnotation:
        trans_anno = TranscriptAnnotation.from_transcript(transcript, region)
        dist1, dist2 = self.pos - region.start + 1, region.end - self.pos + 1
        if transcript.strand == "+":
            trans_anno.na_change = f"c.{cds_size}+{dist1}{self.s_ref}>{self.s_alt}" if dist1 < dist2 else f"c.{cds_size+1}-{dist2}{self.s_ref}>{self.s_alt}"
        else:
            s_ref, s_alt = self.s_ref.reverse_complement(), self.s_alt.reverse_complement()
            trans_anno.na_change = f"c.{transcript.cds_size-cds_size+1}-{dist1}{s_ref}>{s_alt}" if dist1 < dist2 else f"c.{transcript.cds_size-cds_size}+{dist2}{s_ref}>{s_alt}"
        if min(dist1, dist2) <= 2:
            trans_anno.event = "splicing"
            trans_anno.region = "splicing"
        return trans_anno

    def annotate_cds(self, transcript: Transcript, region: Region, cds_size: int) -> TranscriptAnnotation:
        trans_anno = TranscriptAnnotation.from_transcript(transcript, region)
        pos = cds_size + self.pos - region.start + 1
        cdna = transcript.coding_dna
        new_cdna = cdna.substitute(pos, self.s_ref, self.s_alt)
        if transcript.strand == "-":
            cdna = cdna.reverse_complement()
            new_cdna = new_cdna.reverse_complement()
        cstart = cdna.difference_simple(new_cdna)
        na1, na2 = cdna[cstart - 1], new_cdna[cstart - 1]
        trans_anno.na_change = f"c.{cstart}{na1}>{na2}"
        protein = cdna.protein()
        new_protein = new_cdna.protein()
        pstart = protein.difference_simple(new_protein)
        if pstart > protein.size:
            pstart = ((cstart - 1) // 3) + 1
        aa1, aa2 = protein[pstart - 1], new_protein[pstart - 1]
        trans_anno.aa_change = f"p.{aa1}{pstart}{aa2}"
        if aa1 == aa2:
            trans_anno.event = "synonymous"
        else:
            if aa1 == "M" and pstart == 1:
                trans_anno.event = "startloss"
            elif aa1 == "*":
                trans_anno.event = "stoploss"
                trans_anno.aa_change += "ext*?"
            elif aa2 == "*":
                trans_anno.event = "nonsense"
            else:
                trans_anno.event = "missense"
        return trans_anno

    def annotate_mrna(self, transcript: Transcript) -> TranscriptAnnotation:
        region, cds_size, utr_size = transcript.find_region(self.pos)
        if transcript.cds_start <= self.pos <= transcript.cds_end:
            if region.type == RegionType.CDS:
                # CDS
                return self.annotate_cds(transcript, region, cds_size)
            # INTRON
            return self.annotate_intron(transcript, region, cds_size)
        # UTR or UP/DOWN stream
        return self.annotate_utr(transcript, region, utr_size)

    def annotate_ncrna(self, transcript: Transcript) -> TranscriptAnnotation:
        region, _, _ = transcript.find_region(self.pos)
        trans_anno = TranscriptAnnotation.from_transcript(transcript, region)
        trans_anno.region = "ncRNA"
        dna = transcript.dna
        new_dna = dna.substitute(self.pos - transcript.tx_start + 1, self.s_ref, self.s_alt)
        if transcript.strand == "-":
            dna = dna.reverse_complement()
            new_dna = new_dna.reverse_complement()
        start = dna.difference_simple(new_dna)
        na1, na2 = dna[start - 1], new_dna[start - 1]
        trans_anno.na_change = f"n.{start}{na1}>{na2}"
        return trans_anno
